"""
LSTM模型
valid: mae; 2439.35546875， mape: 6.607383728027344, rmse: 2698.279052734375
test: mae; 2685.266845703125， mape: 7.491035461425781, rmse: 3051.925537109375
"""
import os
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import LSTM, Dense, Dropout
from keras.metrics import (RootMeanSquaredError,
                           MeanAbsoluteError,
                           MeanAbsolutePercentageError)


def build_model_old(train_x, nn_nums=50):
    """
    模型构建
    :param train_x:
    :param nn_nums:神经元数量
    :return:
    """
    model = Sequential()
    model.add(LSTM(units=nn_nums,
                   input_shape=(train_x.shape[1], train_x.shape[2]),
                   activation='softsign',
                   return_sequences=True))
    # linear: 激活函数不做任何改变
    model.add(Dropout(rate=0.95))
    model.add(LSTM(units=nn_nums, activation='softsign'))
    model.add(Dropout(rate=0.9))
    model.add(Dense(units=1))

    model.compile(loss='mae', optimizer='RMSProp')
    return model


def build_model(nn_nums=50):
    """
    模型构建
    :param nn_nums:神经元数量
    :return:
    """
    model = Sequential()
    model.add(LSTM(units=nn_nums, activation='relu', dropout=0))
    model.add(Dense(units=1))

    model.compile(loss='mse', optimizer='RMSProp')
    return model


def plot_info(train, test, title, ylabel, pic_path):
    """
    训练 & 测试
    :param train:
    :param test:
    :param title:
    :param ylabel:
    :param pic_path:
    :return:
    """
    plt.clf()

    plt.plot(train)
    plt.plot(test)
    plt.title(title)
    plt.xlabel('Epoch')
    plt.ylabel(ylabel)
    plt.legend(['训练', '测试'], loc='upper left')

    plt.show()
    # pic = '{}/{}.png'.format(pic_path, title)
    # if not os.path.isfile(pic):
    #     plt.savefig(pic)


def lstm_model(data, valid_ratio=0.1, test_ratio=0.1):
    """
    主程序
    :return:
    """
    df = pd.DataFrame(data)
    df.columns = ['data']
    df['forecast'] = df['data'].shift(-1)
    train_data = df.values[: -1]

    scaler = MinMaxScaler(feature_range=(0, 1))
    train_data = scaler.fit_transform(X=train_data)

    train_size = int(len(train_data) * (1 - valid_ratio - test_ratio))
    valid_size = int(len(train_data) * valid_ratio)

    train = train_data[: train_size]
    valid = train_data[train_size: train_size + valid_size]
    test = train_data[train_size + valid_size:]

    train_x = train[:, :-1]
    train_y = train[:, -1:]
    valid_x = valid[:, :-1]
    valid_y = valid[:, -1:]
    test_x = test[:, :-1]
    test_y = test[:, -1:]

    # 每月
    step = 1
    train_x = train_x.reshape((train_x.shape[0], step, train_x.shape[1]))
    train_y = train_y.reshape((train_y.shape[0], step, train_y.shape[1]))
    valid_x = valid_x.reshape((valid_x.shape[0], step, valid_x.shape[1]))
    valid_y = valid_y.reshape((valid_y.shape[0], step, valid_y.shape[1]))
    test_x = test_x.reshape((test_x.shape[0], step, test_x.shape[1]))
    test_y = test_y.reshape((test_y.shape[0], step, test_y.shape[1]))

    model = build_model(nn_nums=40)

    # fit
    history = model.fit(train_x, train_y,
                        epochs=500,
                        batch_size=300,
                        validation_data=(valid_x, valid_y),
                        verbose=1,
                        shuffle=False,
                        workers=8,
                        use_multiprocessing=True)

    loss_metrics = model.evaluate(valid_x, valid_y,
                                  verbose=0,
                                  batch_size=300)

    if isinstance(loss_metrics, list):
        print('loss: {}, metrics: {}'.format(loss_metrics[0], loss_metrics[1:]))
    else:
        print('loss: {}'.format(loss_metrics))

    # loss
    train_loss = history.history['loss']
    val_loss = history.history['val_loss']
    plot_info(train=train_loss, test=val_loss,
              title='LSTM 模型损失值',
              ylabel='Loss', pic_path=None)

    # acc
    if 'acc' in history.history.keys():
        train_acc = history.history['acc']
        test_val_acc = history.history['val_acc']
        plot_info(train=train_acc, test=test_val_acc,
                  title='LSTM 模型准确率',
                  ylabel='Accuracy', pic_path=None)

    pred_train = model.predict(train_x)
    pred_valid = model.predict(valid_x)
    pred_test = model.predict(test_x)

    # 恢复归一化前
    train_y_true = scaler.inverse_transform(train)[:, -1:]
    valid_y_true = scaler.inverse_transform(valid)[:, -1:]
    test_y_true = scaler.inverse_transform(test)[:, -1:]

    pred_train_all = train.copy()
    pred_train_all[:, -1:] = pred_train
    pred_train_scale = scaler.inverse_transform(pred_train_all)
    pred_train_y = pred_train_scale[:, -1:]

    pred_valid_all = valid.copy()
    pred_valid_all[:, -1:] = pred_valid
    pred_valid_scale = scaler.inverse_transform(pred_valid_all)
    pred_valid_y = pred_valid_scale[:, -1:]

    pred_test_all = test.copy()
    pred_test_all[:, -1:] = pred_test
    pred_test_scale = scaler.inverse_transform(pred_test_all)
    pred_test_y = pred_test_scale[:, -1:]

    # 评估
    # mae
    mae = MeanAbsoluteError()
    mae.update_state(y_true=valid_y_true, y_pred=pred_valid_y)
    valid_mae = mae.result()

    mae.reset_states()
    mae.update_state(y_true=test_y_true, y_pred=pred_test_y)
    test_mae = mae.result()

    # mape
    mape = MeanAbsolutePercentageError()
    mape.update_state(y_true=valid_y_true, y_pred=pred_valid_y)
    valid_mape = mape.result()

    mape.reset_states()
    mape.update_state(y_true=test_y_true, y_pred=pred_test_y)
    test_mape = mape.result()

    # rmse
    rmse = RootMeanSquaredError()
    rmse.update_state(y_true=valid_y_true, y_pred=pred_valid_y)
    valid_rmse = rmse.result()

    rmse.reset_states()
    rmse.update_state(y_true=test_y_true, y_pred=pred_test_y)
    test_rmse = rmse.result()

    print('valid: mae; {}， mape: {}, rmse: {}'.
          format(valid_mae, valid_mape, valid_rmse))
    print('test: mae; {}， mape: {}, rmse: {}'.
          format(test_mae, test_mape, test_rmse))

    # 训练，验证，测试画图
    y = [num[0] for num in train_y_true] + \
        [num[0] for num in valid_y_true] + \
        [num[0] for num in test_y_true]

    pred_train_y = [None] + [num[0] for num in pred_train_y] + \
                   [None for _ in pred_valid_y] + \
                   [None for _ in pred_test]
    pred_valid_y = [None for _ in train_y_true] + \
                   [num[0] for num in pred_valid_y] + \
                   [None for _ in pred_test]
    pred_test_y = [None for _ in train_y_true] + \
                  [None for _ in pred_valid] + \
                  [num[0] for num in pred_test_y]
    plt.clf()

    plt.plot(y[: -1], c='b', label='原始数据')
    plt.plot(pred_train_y, c='r', label='训练预测数据')
    plt.plot(pred_valid_y, c='g', label='验证预测数据')
    plt.plot(pred_test_y, c='y', label='测试预测数据')
    plt.legend()

    plt.show()
    raise Exception

    pic = '{}/{}.png'.format(pic_path, model_name + ' 预测效果图')
    if not os.path.isfile(pic):
        plt.savefig(pic)
